package com.gnerv.battle.boot.framework.advice;

import com.gnerv.battle.boot.framework.common.annotation.api.BattleApiResponseEncrypt;
import com.gnerv.battle.boot.framework.common.exception.BattleBusinessException;
import com.gnerv.battle.boot.framework.common.tools.AesTools;
import com.gnerv.battle.boot.framework.common.tools.GsonTools;
import com.gnerv.battle.boot.framework.service.impl.ManageApiCryptoAesService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.Serializable;
import java.nio.charset.StandardCharsets;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

/**
 * 加密响应类
 *
 * @author hermes-di
 * @since 1.0.0.RELEASE
 **/
@Slf4j
@RestControllerAdvice
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice<Object>, Serializable {

    @Resource
    private ManageApiCryptoAesService manageApiCryptoAesService;

    @Override
    public boolean supports(MethodParameter methodParameter, Class<? extends HttpMessageConverter<?>> clazz) {
        return true;
    }

    @Override
    public Object beforeBodyWrite(Object body, MethodParameter methodParameter, MediaType mediaType, Class<? extends HttpMessageConverter<?>> clazz, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        boolean hasMethodAnnotation = methodParameter.hasMethodAnnotation(BattleApiResponseEncrypt.class);
        String encryptUkId = serverHttpRequest.getHeaders().getFirst("Battle-Api-Encrypt");
        if (hasMethodAnnotation) {
            if (StringUtils.hasText(encryptUkId)) {
                try {
                    return encrypt(body, encryptUkId);
                } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException | NoSuchPaddingException |
                         InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
                    throw new RuntimeException(e);
                }
            } else {
                throw new BattleBusinessException("header [Battle-Api-Encrypt] was empty");
            }
        } else if (StringUtils.hasText(encryptUkId)) {
            try {
                return encrypt(body, encryptUkId);
            } catch (NoSuchAlgorithmException | InvalidAlgorithmParameterException | NoSuchPaddingException |
                     InvalidKeyException | IllegalBlockSizeException | BadPaddingException e) {
                throw new RuntimeException(e);
            }
        }
        return body;
    }

    private String encrypt(Object body, String encryptUkId) throws InvalidAlgorithmParameterException, NoSuchPaddingException, IllegalBlockSizeException, NoSuchAlgorithmException, BadPaddingException, InvalidKeyException {
        String json = GsonTools.GSON.toJson(body);
        return encryptAes(json, encryptUkId);
    }

    private String encryptAes(String json, String encryptUkId) throws InvalidAlgorithmParameterException, NoSuchPaddingException, IllegalBlockSizeException, NoSuchAlgorithmException, BadPaddingException, InvalidKeyException {
        String key =  manageApiCryptoAesService.selectKeyByUkId(encryptUkId);
        if (key == null) {
            throw new BattleBusinessException("Secret Key is null");
        }
        return AesTools.encrypt(json, key.getBytes(StandardCharsets.UTF_8));
    }

}